"""
泰坦尼克号幸存者预测项目主程序
"""

import os
import sys
import time
import logging
import warnings
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import joblib
from datetime import datetime

# 导入自定义模块
from src.data_preprocessing import load_data, explore_data, handle_missing_values, split_data
from src.feature_engineering import apply_feature_engineering_pipeline
from src.model_training import train_and_select_best_model, save_model, load_model
from src.model_evaluation import evaluate_model, plot_confusion_matrix, plot_feature_importance, analyze_errors, plot_roc_curve
from src.visualization import setup_visualization_style, plot_survival_distribution, plot_correlation_heatmap, create_eda_dashboard

# 忽略警告
warnings.filterwarnings('ignore')

# 设置日志
def setup_logging():
    """设置日志记录"""    
    # 创建logs目录
    logs_dir = 'logs'
    os.makedirs(logs_dir, exist_ok=True)
    
    # 生成日志文件名（包含时间戳）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(logs_dir, f'titanic_{timestamp}.log')
    
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)
        ]
    )
    
    return logging.getLogger(__name__)


def check_data_files(data_dir):
    """检查数据文件是否存在
    
    参数:
        data_dir (str): 数据目录路径
        
    返回:
        tuple: (train_path, test_path) 数据文件路径
        
    异常:
        FileNotFoundError: 如果数据文件不存在
    """
    train_path = os.path.join(data_dir, 'train.csv')
    test_path = os.path.join(data_dir, 'test.csv')
    
    if not os.path.exists(train_path):
        raise FileNotFoundError(f"训练数据文件不存在: {train_path}\n请从Kaggle下载泰坦尼克号数据集并放在data目录下")
    
    if not os.path.exists(test_path):
        raise FileNotFoundError(f"测试数据文件不存在: {test_path}\n请从Kaggle下载泰坦尼克号数据集并放在data目录下")
    
    return train_path, test_path


def create_submission_file(model, test_df, encoders, scaler, submission_path):
    """创建提交文件
    
    参数:
        model: 训练好的模型
        test_df (pd.DataFrame): 测试数据
        encoders (dict): 特征编码器
        scaler (StandardScaler): 特征标准化器
        submission_path (str): 提交文件路径
    """
    logger = logging.getLogger(__name__)
    logger.info("创建预测结果提交文件...")
    
    try:
        # 处理测试数据
        test_processed = handle_missing_values(test_df)
        
        # 应用特征工程
        test_fe, _, _ = apply_feature_engineering_pipeline(
            test_processed, is_train=False, encoders=encoders, scaler=scaler
        )
        
        # 确保特征顺序与训练集一致
        # 这里需要从模型中获取特征名称，但不同模型可能有不同方式
        # 简化处理：假设测试集特征与训练集特征顺序一致
        
        # 进行预测
        predictions = model.predict(test_fe)
        
        # 创建提交DataFrame
        submission = pd.DataFrame({
            'PassengerId': test_df['PassengerId'],
            'Survived': predictions.astype(int)
        })
        
        # 保存提交文件
        submission.to_csv(submission_path, index=False)
        
        logger.info(f"提交文件已保存到: {submission_path}")
        logger.info(f"预测结果分布: 幸存={predictions.sum()} ({predictions.sum()/len(predictions)*100:.2f}%), "
                   f"遇难={(1-predictions).sum()} ({(1-predictions).sum()/len(predictions)*100:.2f}%)")
        
        return submission
    except Exception as e:
        logger.error(f"创建提交文件时出错: {e}")
        raise


def main(args=None):
    """
    主函数，运行整个项目流程
    
    参数:
        args (dict): 命令行参数或配置参数
    """
    # 设置日志
    logger = setup_logging()
    
    # 记录开始时间
    start_time = time.time()
    logger.info("泰坦尼克号幸存者预测项目开始")
    
    # 默认参数
    if args is None:
        args = {
            'data_dir': 'data',
            'models_dir': 'models',
            'optimize': True,
            'create_submission': True,
            'save_plots': True,
            'encoding_method': 'label'  # 'label' 或 'onehot'
        }
    
    try:
        # 设置可视化风格
        setup_visualization_style()
        
        # 确保目录存在
        os.makedirs(args['data_dir'], exist_ok=True)
        os.makedirs(args['models_dir'], exist_ok=True)
        os.makedirs('notebooks', exist_ok=True)
        os.makedirs('logs', exist_ok=True)
        os.makedirs('results', exist_ok=True)
        
        # 检查数据文件
        train_path, test_path = check_data_files(args['data_dir'])
        
        # 1. 数据准备阶段
        logger.info("=" * 50)
        logger.info("1. 数据准备阶段")
        logger.info("=" * 50)
        
        # 加载数据
        train_df, test_df = load_data(train_path, test_path)
        
        # 探索数据
        logger.info("训练数据探索:")
        train_info = explore_data(train_df)
        
        # 创建EDA仪表板并保存
        if args['save_plots']:
            try:
                eda_fig = create_eda_dashboard(train_df)
                eda_path = os.path.join('results', 'eda_dashboard.png')
                eda_fig.savefig(eda_path, dpi=300, bbox_inches='tight')
                logger.info(f"EDA仪表板已保存到: {eda_path}")
                plt.close(eda_fig)
            except Exception as e:
                logger.warning(f"创建EDA仪表板时出错: {e}")
        
        # 2. 数据预处理阶段
        logger.info("=" * 50)
        logger.info("2. 数据预处理阶段")
        logger.info("=" * 50)
        
        # 处理缺失值
        train_df_processed = handle_missing_values(train_df)
        
        # 3. 特征工程
        logger.info("=" * 50)
        logger.info("3. 特征工程")
        logger.info("=" * 50)
        
        # 应用特征工程流程
        train_df_fe, encoders, scaler = apply_feature_engineering_pipeline(
            train_df_processed, is_train=True, encoding_method=args['encoding_method']
        )
        
        # 4. 数据集划分
        logger.info("=" * 50)
        logger.info("4. 数据集划分")
        logger.info("=" * 50)
        
        # 分割数据集
        X_train, X_val, y_train, y_val = split_data(train_df_fe)
        
        # 5. 模型训练
        logger.info("=" * 50)
        logger.info("5. 模型训练")
        logger.info("=" * 50)
        
        # 训练并选择最佳模型
        best_model, evaluations = train_and_select_best_model(
            X_train, y_train, X_val, y_val, optimize=args['optimize']
        )
        
        # 6. 模型评估
        logger.info("=" * 50)
        logger.info("6. 模型评估")
        logger.info("=" * 50)
        
        # 在验证集上评估最佳模型
        val_metrics = evaluate_model(best_model, X_val, y_val)
        
        # 绘制混淆矩阵
        y_val_pred = best_model.predict(X_val)
        cm_fig = plot_confusion_matrix(y_val, y_val_pred, title="Confusion Matrix")
        if args['save_plots']:
            cm_path = os.path.join(args['models_dir'], 'confusion_matrix.png')
            cm_fig.savefig(cm_path, dpi=300, bbox_inches='tight')
            logger.info(f"混淆矩阵已保存到: {cm_path}")
        plt.close(cm_fig)
        
        # 绘制ROC曲线（如果模型支持概率预测）
        if hasattr(best_model, 'predict_proba'):
            try:
                y_val_proba = best_model.predict_proba(X_val)[:, 1]
                roc_fig = plot_roc_curve(y_val, y_val_proba, title="ROC Curve")
                if args['save_plots']:
                    roc_path = os.path.join(args['models_dir'], 'roc_curve.png')
                    roc_fig.savefig(roc_path, dpi=300, bbox_inches='tight')
                    logger.info(f"ROC曲线已保存到: {roc_path}")
                plt.close(roc_fig)
            except Exception as e:
                logger.warning(f"绘制ROC曲线时出错: {e}")
        
        # 绘制特征重要性（仅适用于有feature_importances_属性的模型）
        if hasattr(best_model, 'feature_importances_'):
            fi_fig = plot_feature_importance(best_model, X_train.columns, title="Feature Importance")
            if args['save_plots']:
                fi_path = os.path.join(args['models_dir'], 'feature_importance.png')
                fi_fig.savefig(fi_path, dpi=300, bbox_inches='tight')
                logger.info(f"特征重要性图已保存到: {fi_path}")
            plt.close(fi_fig)
        
        # 分析错误预测
        error_samples, error_analysis = analyze_errors(best_model, X_val, y_val, X_train.columns)
        
        # 7. 保存模型
        logger.info("=" * 50)
        logger.info("7. 保存模型")
        logger.info("=" * 50)
        
        # 保存最佳模型
        model_path = os.path.join(args['models_dir'], 'best_model.pkl')
        save_model(best_model, model_path)
        
        # 保存编码器和标准化器
        joblib.dump(encoders, os.path.join(args['models_dir'], 'encoders.pkl'))
        joblib.dump(scaler, os.path.join(args['models_dir'], 'scaler.pkl'))
        
        # 保存评估结果
        results = {
            'model_type': type(best_model).__name__,
            'validation_metrics': val_metrics,
            'evaluations': evaluations,
            'feature_importance': dict(zip(X_train.columns, best_model.feature_importances_)) if hasattr(best_model, 'feature_importances_') else None,
            'error_analysis': error_analysis
        }
        
        results_path = os.path.join('results', 'evaluation_results.pkl')
        joblib.dump(results, results_path)
        logger.info(f"评估结果已保存到: {results_path}")
        
        # 8. 创建提交文件（可选）
        if args['create_submission']:
            logger.info("=" * 50)
            logger.info("8. 创建提交文件")
            logger.info("=" * 50)
            
            submission_path = os.path.join(args['data_dir'], 'submission.csv')
            submission = create_submission_file(best_model, test_df, encoders, scaler, submission_path)
        
        # 计算并记录总运行时间
        end_time = time.time()
        elapsed_time = end_time - start_time
        logger.info("=" * 50)
        logger.info("项目完成!")
        logger.info("=" * 50)
        logger.info(f"最佳模型: {type(best_model).__name__}")
        logger.info(f"验证集准确率: {val_metrics['accuracy']:.4f}")
        logger.info(f"验证集F1分数: {val_metrics['f1']:.4f}")
        logger.info(f"总运行时间: {elapsed_time:.2f} 秒")
        
        return {
            'model': best_model,
            'metrics': val_metrics,
            'encoders': encoders,
            'scaler': scaler,
            'elapsed_time': elapsed_time
        }
        
    except Exception as e:
        logger.error(f"项目运行过程中出错: {e}", exc_info=True)
        raise


def parse_arguments():
    """解析命令行参数（简化版）
    
    返回:
        dict: 解析后的参数
    """
    # 这里简化处理，实际应用中可以使用argparse等库
    args = {
        'data_dir': 'data',
        'models_dir': 'models',
        'optimize': True,
        'create_submission': True,
        'save_plots': True,
        'encoding_method': 'label'  # 'label' 或 'onehot'
    }
    
    # 这里可以添加从命令行解析参数的代码
    # 例如:
    # import argparse
    # parser = argparse.ArgumentParser(description='泰坦尼克号幸存者预测')
    # parser.add_argument('--data_dir', type=str, default='data', help='数据目录')
    # ...其他参数
    # args = vars(parser.parse_args())
    
    return args


if __name__ == "__main__":
    # 解析命令行参数
    args = parse_arguments()
    
    # 运行主程序
    try:
        results = main(args)
        print("\n程序成功执行完成!")
    except Exception as e:
        print(f"\n程序执行出错: {e}")
        sys.exit(1)
